import os
import torch
import numpy as np
import random
import time
from src.data.dataloader import get_har_dataloader
from src.models.attention_model import HARAttention

def test_single_random_sample(model_path='models/attention/best_model.pth', device='cpu'):
    """
    随机选择一个测试样本并进行预测，直接从文件加载数据
    
    参数:
        model_path: 模型权重文件路径
        device: 使用的设备 ('cpu' 或 'cuda')
    """
    # 使用当前时间作为随机种子，确保每次运行都不同
    current_time = int(time.time())
    random.seed(current_time)
    print(f"使用随机种子: {current_time}")
    
    # 设置数据集路径
    data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data', 'har_dataset')
    
    # 加载活动标签
    activity_labels = {}
    labels_file = os.path.join(data_dir, 'activity_labels.txt')
    if os.path.exists(labels_file):
        with open(labels_file, 'r') as f:
            for line in f:
                idx, label = line.strip().split(' ', 1)
                activity_labels[int(idx) - 1] = label  # 减1是因为我们的模型输出从0开始
    else:
        print(f"警告: 活动标签文件不存在: {labels_file}")
        return
    
    # 直接从文件加载数据
    X_test_file = os.path.join(data_dir, 'test', 'X_test.txt')
    y_test_file = os.path.join(data_dir, 'test', 'y_test.txt')
    
    if not (os.path.exists(X_test_file) and os.path.exists(y_test_file)):
        print(f"错误: 测试数据文件不存在")
        return
    
    # 计算测试集大小
    with open(y_test_file, 'r') as f:
        test_size = sum(1 for _ in f)
    
    # 随机选择一个样本索引
    random_idx = random.randint(0, test_size - 1)
    print(f"随机选择样本索引: {random_idx}")
    
    # 加载特征
    X_test = None
    with open(X_test_file, 'r') as f:
        for i, line in enumerate(f):
            if i == random_idx:
                try:
                    values = [float(x) for x in line.strip().split()]
                    X_test = values
                    break
                except Exception as e:
                    print(f"解析特征数据时出错: {e}")
                    return
    
    # 加载标签
    y_test = None
    with open(y_test_file, 'r') as f:
        for i, line in enumerate(f):
            if i == random_idx:
                try:
                    y_test = int(line.strip()) - 1  # 减1是因为标签从1开始
                    break
                except Exception as e:
                    print(f"解析标签数据时出错: {e}")
                    return
    
    if X_test is None or y_test is None:
        print("错误: 无法加载样本数据")
        return
    
    # 将特征转换为张量
    features = torch.tensor([X_test], dtype=torch.float32).to(device)
    
    # 打印特征的基本统计信息
    print(f"\n特征统计信息:")
    print(f"特征形状: {features.shape}")
    print(f"特征均值: {features.mean().item():.4f}")
    print(f"特征标准差: {features.std().item():.4f}")
    print(f"特征最小值: {features.min().item():.4f}")
    print(f"特征最大值: {features.max().item():.4f}")
    
    # 加载模型
    input_size = features.shape[1]
    num_classes = len(activity_labels)
    
    # 创建模型实例
    model = HARAttention(input_size=input_size, num_classes=num_classes)
    
    # 检查模型权重文件是否存在
    if not os.path.exists(model_path):
        print(f"错误: 模型文件 {model_path} 不存在!")
        return
    
    # 加载预训练权重
    try:
        state_dict = torch.load(model_path, map_location=device)
        model.load_state_dict(state_dict)
        print("模型权重加载成功!")
    except Exception as e:
        print(f"模型权重加载失败: {e}")
        return
    
    model.to(device)
    model.eval()
    
    # 进行预测
    with torch.no_grad():
        outputs = model(features)
        _, predicted = torch.max(outputs, 1)
        predicted_label = predicted.item()
    
    # 打印结果
    print("\n===== 人类活动识别测试 =====")
    print(f"真实活动: {activity_labels[y_test]} (类别 {y_test})")
    print(f"预测活动: {activity_labels[predicted_label]} (类别 {predicted_label})")
    
    # 打印原始输出
    print("\n模型原始输出:")
    for i, out in enumerate(outputs[0]):
        print(f"类别 {i} ({activity_labels[i]}): {out.item():.4f}")
    
    # 计算所有类别的概率
    probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
    
    # 打印每个类别的概率
    print("\n各活动类别的概率:")
    for i, prob in enumerate(probabilities):
        print(f"{activity_labels[i]}: {prob.item():.4f} ({prob.item()*100:.2f}%)")
    
    # 判断预测是否正确
    if predicted_label == y_test:
        print("\n✓ 预测正确!")
    else:
        print("\n✗ 预测错误!")

if __name__ == "__main__":
    # 检查是否有可用的CUDA设备
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")
    
    # 指定模型路径，默认使用attention模型
    model_path = 'models/attention/best_model.pth'
    
    # 运行测试
    test_single_random_sample(model_path, device)